#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/5/29 10:16
# @Author  : liutianwei
# @File    : HawkDecoderService.py
# @Software: PyCharm
import pdb
import json
import os
import re
import threading
import time
from ctypes import *
from conftest import logging
import allure
from src.utils.common import method, async_color_print, get_timestamp, pop, uuid
from src.core.Status.HawkDecoderStatus import SpeechEngineStatusCode, SpeechEngineParam
from src.core.Common.Bean.RuntimeRecorder import RuntimeRecorder
from src.core.Status.HawkDecoderStatus import MongoCode


# 定义回调函数结构体
class SpeechEngineResultSt(Structure):
    _fields_ = [
        ("time", c_ulong),
        ("flag", c_int),
        ("back", c_int),
        ("data", c_char_p)
    ]


class HawkDecoderClient:
    def __init__(self, lib_path: str):
        self.default_language = None
        self.m_status = None
        self.m_caseList = None
        self.m_config = None
        self.m_library = None
        self.m_hawkEngine = None
        self.m_callback = None
        self.m_index = None
        self.m_recorder = None
        self.m_sampleRate = 16000
        self.speech_end_time = 0
        self.m_delay = 1.0
        self.m_personal_tag = 0
        self.m_lock = threading.Lock()
        self.m_create_cond = threading.Condition(self.m_lock)
        self.m_start_cond = threading.Condition(self.m_lock)
        self.m_finish_cond = threading.Condition(self.m_lock)
        self.m_cancel_time = -1
        self.m_cancel_event = None
        self.m_process_timescale_callback = None
        self.m_region_list = []
        try:
            self.m_library = cdll.LoadLibrary(lib_path)
        except Exception as e:
            print(f"load HawkDecoder library error: {e}")

    @method(CFUNCTYPE(c_int, c_int, POINTER(SpeechEngineResultSt)))
    def callback_func(self, status: int, result: SpeechEngineResultSt):
        call_data = result[0].data
        if call_data is None or len(call_data) == 0:
            return 0
        _data = call_data.decode(encoding="utf-8", errors='ignore')
        if status in [SpeechEngineStatusCode.STATUS_SPEECH_SR_INIT_SUCCESS, SpeechEngineStatusCode.STATUS_SPEECH_SR_INIT_FAILED]:
            with self.m_lock:
                self.m_status = status
                self.m_create_cond.notify()

        if status in [SpeechEngineStatusCode.STATUS_SPEECH_SR_START_SUCCESS, SpeechEngineStatusCode.STATUS_SPEECH_SR_START_FAILED]:
            with self.m_lock:
                self.m_recorder.process_start_time = time.time()
                self.m_status = status
                async_color_print((f"--------------------- Session Start [{self.m_index}] --------------------", "yellow"))
                self.m_start_cond.notify()
        
        self.notify_result(status, _data)

        if status is SpeechEngineStatusCode.STATUS_SPEECH_SESSION_FINISH:
            with self.m_lock:
                self.m_recorder.process_end_time = time.time()
                self.m_status = status
                self.m_finish_cond.notify()
                self.m_recorder.cpuTime = (self.m_recorder.process_end_time - self.m_recorder.process_start_time) * 1000
                if self.m_recorder.voiceTime:
                    self.m_recorder.rt = self.m_recorder.cpuTime / self.m_recorder.voiceTime
                else:
                    self.m_recorder.rt = 0.0
                voice_length = self.m_recorder.voiceTime / 1000.0
                cpu_time = self.m_recorder.cpuTime / 1000.0
                rt = self.m_recorder.rt
                async_color_print(("[voice length:{:.3f}(s)][cpu time:{:.3f}(s)][rt:{:.3f}]".format(voice_length, cpu_time, rt), None))
                async_color_print((f"-------------------- Session Finish [{self.m_index}] --------------------", "yellow"))
                
                stop_tag = 'inf'
                if self.m_region_list:
                    stop_tag = pop(self.m_region_list,0)
                _data = {"source": "Mongo", "type": "ASSERTED", "voiceRegionEnd": float(stop_tag)}
                self.notify_result(MongoCode.STATUS_MONGO_FINISH, _data)

        return 0

    def attach(self, config: str, log_path) -> int:
        if not config:
            logging.error("config is empty")
            return -1

        self.m_hawkEngine = self.m_library.attach_speech_engine(config.encode("utf-8"), None, self.callback_func, None)
        self.set_speech_engine_log_path(log_path)
        if self.m_hawkEngine is None:
            logging.error(f"create HawkDecoder engine failed the engine is:{self.m_hawkEngine}")
            return -2

        with self.m_lock:
            self.m_create_cond.wait(timeout=60)
            if self.m_status == SpeechEngineStatusCode.STATUS_SPEECH_SR_INIT_SUCCESS:
                logging.info('Create HAWK speech engine success!')
                return 0
            else:
                logging.error(f'Create HAWK speech engine failed, status:{self.m_status}, engine:{self.m_hawkEngine}')
                return -1

    def start(self, index, channel):
        with self.m_lock:
            self.m_index = index
            logging.info("调用start接口...")
            ret = self.m_library.start_speech_engine(self.m_hawkEngine, channel, None)
            self.m_start_cond.wait(timeout=5)
            if self.m_status == SpeechEngineStatusCode.STATUS_SPEECH_SR_START_SUCCESS:
                logging.info(f'start HAWK speech engine success!')
                return 0
            elif self.m_status == SpeechEngineStatusCode.STATUS_SPEECH_SR_START_FAILED:
                logging.error(f'start HAWK speech engine failed ret:{ret}')
                return -1
            else:
                logging.error(f'start HAWK speech engine failed ret:{ret}')
                return -1

    def stop(self):
        with self.m_lock:
            ret = self.m_library.stop_speech_engine(self.m_hawkEngine)
            if ret != 0:
                return -1
            self.m_finish_cond.wait(timeout=20)
            logging.info('Stop aibs engine stopped successfully!')

    def set_speech_engine_log_path(self, log_path: str):
        ret = self.m_library.set_speech_engine_log_path(self.m_hawkEngine, log_path.encode(encoding='utf-8'))
        if ret != 0:
            print(f"set speech engine log error: {ret}")
            return -1
        return 0

    def speech_engine_set_language_info(self, lang: str) -> int:
        # 如果之前的Case已经设置过语种了，并且和当前的语种相等，则跳过设置语种的逻辑
        if self.default_language is not None and self.default_language == lang:
            return 0

        try:
            if lang not in ["cmn", "eng", "yue", "chuan"]:
                raise ValueError
        except ValueError:
            logging.error(f"Unknown lang: {lang}.")
            return -1
        lang_json = {"wakeup":[lang],"asr":[lang]}
        json_data = json.dumps(lang_json, ensure_ascii=False, separators=(',', ':'))
        ret = self.m_library.speech_engine_set_language_info(self.m_hawkEngine, json_data.encode('utf-8'))
        if ret != 0:
            logging.error(f'set Hawk speech engine lang error ret: {ret}')
            return ret

        self.default_language = lang
        return 0

    def add_wakeup_word(self, wakeupAils, threshold) -> int:
        try:
            if not wakeupAils:
                raise ValueError
            threshold = float(threshold)
        except ValueError:
            logging.error(f"Set alis wakeup word error: {wakeupAils}, {threshold}.")
            return -1
        ret = self.m_library.add_wakeup_word(self.m_hawkEngine, wakeupAils.encode(encoding='utf-8'), c_float(threshold))
        if ret != 0:
            logging.error(f'set Hawk speech engine lang error ret: {ret}')
            return ret
        return 0

    def get_wakeup_term(self):
        return self.m_library.get_wakeup_term(self.m_hawkEngine)
 
    def set_speech_engine_param(self, param, value):
        if type(value) is int:
            ret = self.m_library.set_speech_engine_param(self.m_hawkEngine, param, byref(c_uint(value)))
        elif type(value) is float:
            ret = self.m_library.set_speech_engine_param(self.m_hawkEngine, param, byref(c_float(value)))
        elif type(value) is str:
            ret = self.m_library.set_speech_engine_param(self.m_hawkEngine, param, byref(c_char_p(value)))
        elif type(value) is bool:
            ret = self.m_library.set_speech_engine_param(self.m_hawkEngine, param, byref(c_bool(value)))
        else:
            logging.error(f'Set HAWK speech engine parm error, unknown type of value: {type(value)}')
            return -1
        if ret != 0:
            logging.error(f'Set HAWK speech engine parm error, ret: {ret}')
            return ret
        return 0

    def set_speech_engine_work_mode(self, mode):
        try:
            mode = int(mode)
            if mode not in [0, 1, 2]:
                raise ValueError
        except ValueError:
            logging.error(f"Unknown work mode：{mode}.")
            return -1
        ret = self.m_library.set_speech_engine_work_mode(self.m_hawkEngine, c_int(mode))
        if ret != 0:
            logging.error(f'set Hawk speech engine work mode error ret: {ret}')
            return ret
        return 0

    def speech_engine_set_data_type(self, data_type):
        """
        设置处理音频数据类型
        :param data_type: 数据类型, 1:数据不需要降噪; 2:数据需要降噪; 0:单声道数据,不需要降噪
        :return: 接口返回状态码
        """
        try:
            data_type = int(data_type)
            if data_type not in [0, 1, 2]:
                raise ValueError
        except ValueError:
            logging.error(f"Unknown data_type = {data_type}")
            return -1
        self.m_library.speech_engine_set_data_type(self.m_hawkEngine, c_int(data_type))
        return 0

    def cancel(self, flag: int):
        self.m_recorder.cancel_start_time = time.time()
        ret = self.m_library.cancel_speech_engine(self.m_hawkEngine, flag)
        self.m_recorder.cancel_end_time = time.time()
        if ret != 0:
            logging.error(f'cancel HAWK speech engine error ret: {ret}')
        return 0

    def free(self):
        return self.m_library.free_speech_engine(self.m_hawkEngine)

    def set_recorder(self, recoder: RuntimeRecorder):
        self.m_recorder = recoder

    # def update_personalized_info(self):
    #     ret = self.m_library.update_personalized_info(self.m_hawkEngine)

    def speech_engine_open_voice_input(self, param):
        return self.m_library.speech_engine_open_voice_input(self.m_hawkEngine, param)

    def speech_engine_close_voice_input(self):
        return self.m_library.speech_engine_close_voice_input(self.m_hawkEngine)

    def set_wakeup_word_enable(self, word: str, option):
        try:
            option = int(option)
            if (option not in [0, 1]) or not word:
                raise ValueError
            option = bool(option)
        except ValueError:
            logging.error(f"set wakeup word enable error, nknown word: {word} or option = {option}")
            return -1
        ret = self.m_library.set_wakeup_word_enable(self.m_hawkEngine, word.encode(encoding='utf-8'), c_bool(option))
        if ret != 0:
            logging.error(f'set Hawk speech engine wakeup word enable error ret: {ret}')
            return ret
        return 0

    def set_tts(self, state: bool):
        return self.m_library.set_tts_state(self.m_hawkEngine, state)

    def set_hmi(self, file: str) -> int:
        if not file:
            return -1
        try:
            if not os.path.exists(file):
                raise ValueError
            with open(file, 'r', encoding='utf-8') as fp:
                dict_data = json.load(fp)

            et_json = {"et_id": uuid(), "mts": get_timestamp()}
            dict_data.update(et_json)
            json_data = json.dumps(dict_data, ensure_ascii=False, separators=(',', ':'))
            with allure.step(f"设置HMI信息: {dict_data.get('type')}"):
                with allure.step(f"{json_data}"):
                    pass
            return self.m_library.set_speech_engine_param(SpeechEngineParam.SPEECH_ENGINE_PARAM_HMI_CONTEXT, json_data.encode('utf-8'))
        except Exception as e:
            logging.error(f"Set SpeechEngin hmi error, parse json data error:{e}")
            return -1
    
    def set_listener(self, callback):
        self.m_callback = callback

    def notify_result(self, status, callback_data):
        if self.m_callback:
            self.m_callback(status, callback_data)

    def process_data(self, audioPath: str, frameSize: int, timeMode: bool=False) -> int:
        """
        送单声道音频
        :param audioPath: 音频路径
        :param frameSize: 每一步的读取size
        :param timeMode: 送音频方式
        :return:
        """
        if (not audioPath) or frameSize == 0:
            logging.error(f'Process data error, audio_path:{audioPath}, step_size:{frameSize}')
            return -2

        channelNum = int(frameSize / 320)
        audio_length = os.path.getsize(audioPath)

        timestamp = get_timestamp() if timeMode else 0

        audio_timestamp = 0

        stop_tag = None
        if self.m_region_list:
            stop_tag = self.m_region_list[0]

        self.start_record()
        # 循环读取，并送音频
        with open(audioPath, 'rb') as fp:

            # 获取文件头信息
            header = fp.read(100)
            if header[:4] == b"RIFF":
                # 如果是wav文件，则获取头信息中的音频长度
                pos = re.search(b'data', header).end()
                audio_length = int.from_bytes(header[pos:pos + 4], byteorder='little')
                # 将文件指针移动到音频数据的起始位置
                fp.seek(pos + 4, 0)
            else:
                fp.seek(0, 0)
            # 实际pcm音频流数据
            remaining_bytes = audio_length
            self.m_recorder.fileSize = remaining_bytes
            self.m_recorder.voiceTime = 1000 * self.m_recorder.fileSize / (2.0 * self.m_sampleRate * channelNum)

            while remaining_bytes > 0:
                if remaining_bytes >= frameSize:
                    # 读取完整一帧数据
                    data = fp.read(frameSize)
                else:
                    # 读取剩余的不足一帧的数据
                    data = fp.read(remaining_bytes)

                if len(data) == 0:
                    break
                
                timestamp += 10
                audio_timestamp += 10
                ret = self.m_library.speech_engine_process_data(self.m_hawkEngine, data, len(data), c_ulong(timestamp))
                # if ret != 0:
                #     if self.m_status == SpeechEngineStatusCode.STATUS_SPEECH_SESSION_FINISH:
                #         fp.close()
                #         return 0
                #     logging.error("Process data failed or the current task don't need the rest of data.")
                #     print("Process data failed or the current task don't need the rest of data.")
                #     return -1

                # 每送一包，把时间戳callback出去
                if self.m_process_timescale_callback:
                    self.m_process_timescale_callback(audio_timestamp)

                # 如果存在Delay，speed=1时，每帧delay10ms（0.01s）
                delay_time = 0.01 * self.m_delay
                time.sleep(delay_time)

                remaining_bytes -= len(data)
                # 如果已经读取了指定长度，则退出循环
                if remaining_bytes <= 0:
                    break
        self.stop_record()
        return 0

    def process_delay_data(self, audioPath: str, frameSize: int,  time_list: list, timeMode: bool = False):
        """
        送单声道音频
        :param audioPath: 音频路径
        :param frameSize: 每一步的读取size
        :param timeMode: 送音频方式
        :return:
        """
        expect_local_time = []
        if (not audioPath) or frameSize == 0:
            logging.error(f'Process data error, audio_path:{audioPath}, step_size:{frameSize}')
            return -2

        channelNum = int(frameSize / 320)
        audio_length = os.path.getsize(audioPath)

        timestamp = get_timestamp() if timeMode else 0

        stop_tag = None
        if self.m_region_list:
            stop_tag = self.m_region_list[0]

        self.start_record()
        # 循环读取，并送音频
        with open(audioPath, 'rb') as fp:

            # 获取文件头信息
            header = fp.read(100)
            if header[:4] == b"RIFF":
                # 如果是wav文件，则获取头信息中的音频长度
                pos = re.search(b'data', header).end()
                audio_length = int.from_bytes(header[pos:pos + 4], byteorder='little')
                # 将文件指针移动到音频数据的起始位置
                fp.seek(pos + 4, 0)
            else:
                fp.seek(0, 0)
            # 实际pcm音频流数据
            remaining_bytes = audio_length
            self.m_recorder.fileSize = remaining_bytes
            self.m_recorder.voiceTime = 1000 * self.m_recorder.fileSize / (2.0 * self.m_sampleRate * channelNum)

            while remaining_bytes > 0:
                if remaining_bytes >= frameSize:
                    # 读取完整一帧数据
                    data = fp.read(frameSize)
                else:
                    # 读取剩余的不足一帧的数据
                    data = fp.read(remaining_bytes)

                if len(data) == 0:
                    break

                timestamp += 10
                if timestamp in time_list:
                    expect_local_time.append({f'{timestamp}': get_timestamp()})
                ret = self.m_library.speech_engine_process_data(self.m_hawkEngine, data, len(data), c_ulong(timestamp))
                # if ret != 0:
                #     if self.m_status == SpeechEngineStatusCode.STATUS_SPEECH_SESSION_FINISH:
                #         fp.close()
                #         return 0
                #     logging.error("Process data failed or the current task don't need the rest of data.")
                #     print("Process data failed or the current task don't need the rest of data.")
                #     return -1

                # 每送一包，把时间戳callback出去
                if self.m_process_timescale_callback:
                    self.m_process_timescale_callback(timestamp)

                # 如果存在Delay，speed=1时，每帧delay10ms（0.01s）
                delay_time = 0.01 * self.m_delay
                time.sleep(delay_time)

                remaining_bytes -= len(data)
                # 如果已经读取了指定长度，则退出循环
                if remaining_bytes <= 0:
                    break
        self.stop_record()
        return expect_local_time

    def update_personalized_info(self, file, option):
        self.m_library.update_personalized_info.argtypes = [c_void_p, c_char_p, c_char_p, c_float]
        self.m_library.update_personalized_info.restype = c_int
        self.m_personal_tag = option
        if not self.m_personal_tag:
            return 0
        if self.m_status == SpeechEngineStatusCode.STATUS_SPEECH_SR_INIT_FAILED:
            logging.error('Aibs engine created failed, please checkout before set personalized info!')
            return -1
        if not file:
            logging.error(f'Aibs engine personal tag list is none: {file}')
            return -1

        with open(file, "r", encoding="utf-8") as fp:
            if fp is None:
                logging.error(f"Set personalized info error, open tag list failed:{file}")
                return -2
            while True:
                line = fp.readline().strip()
                if not line:
                    break
                self.m_personal_tag = 0
                self.m_lock.acquire()
                try:
                    jsons = json.dumps(json.loads(line), ensure_ascii=False)
                except Exception as e:
                    print(f"parse json file [{file}] error: {e}.")
                    self.m_lock.release()
                    fp.close()
                    return -1

                self.m_recorder.personalized_start_time = time.time()
                ret = self.m_library.update_personalized_info(self.m_hawkEngine, jsons.encode('utf-8'), "utf-8".encode(), c_float(0.5))
                if ret != 0:
                    fp.close()
                    self.m_lock.release()
                    return -1

                self.m_recorder.personalized_end_time = time.time()

                # TODO 需要在回调json中解析消息并赋值self.m_personal_tag = 48
                if self.m_personal_tag == 48:
                    self.m_lock.release()
                    self.m_recorder.personalized_load_time = self.m_recorder.personalized_end_time - self.m_recorder.personalized_start_time
                    print("m_recorder->tagLoadTime:%d\n", self.m_recorder.personalized_load_time)
                    fp.close()
                    return 0

                # TODO 需要在回调json中解析消息并赋值self.m_personal_tag = 49
                elif self.m_personal_tag == 49:
                    self.m_lock.release()
                    logging.error(f"Set personal tag list failed:{file}")
                    fp.close()
                    return -1
                else:
                    self.m_lock.release()
                    fp.close()
                    return 0
        fp.close()
        return 0

    def set_freetalk_state(self, switch, channel):
        ret = self.m_library.set_freetalk_state(self.m_hawkEngine, byref(c_bool(switch)), channel)
        if ret != 0:
            logging.error(f"Set freetalk state error ret; {ret}")
            return -1
        return 0

    def start_record(self):
        self.m_library.speech_start_record.argtypes = [c_void_p]
        self.m_library.speech_start_record.restype = c_int
        return self.m_library.speech_start_record(self.m_hawkEngine)

    def stop_record(self):
        self.m_library.speech_stop_record.argtypes = [c_void_p]
        self.m_library.speech_stop_record.restype = c_int
        return self.m_library.speech_stop_record(self.m_hawkEngine)

    def set_delay(self, delay: float):
        self.m_delay = delay

    def get_enroll_text(self):
        self.m_library.aibs_get_enroll_text.argtypes = [c_void_p]
        self.m_library.aibs_get_enroll_text.restype = c_char_p
        result =  self.m_library.aibs_get_enroll_text(self.m_engine)
        return result.decode(encoding="utf-8")

    def aibs_sensitive_word_judgment(self, word):
        self.m_library.aibs_sensitive_word_judgment.argtypes = [c_void_p, c_char_p]
        self.m_library.aibs_sensitive_word_judgment.restype = c_char_p
        result =  self.m_library.aibs_sensitive_word_judgment(self.m_engine, word.encode())
        return result

    def aibs_get_registered_speaker(self):
        self.m_library.aibs_get_registered_speaker.argtypes = [c_void_p]
        self.m_library.aibs_get_registered_speaker.restype = c_char_p
        result =  self.m_library.aibs_get_registered_speaker(self.m_engine)
        return result.decode(encoding="utf-8")
    
    def start_enroll(self, user_id, text, index, channel):
        buffer = {"type":"startEnroll","data":{"channel":channel,"id":index,"user_id":user_id,"text":text,"start_time":-1,"end time":-1}}
        self.m_library.set_speech_engine_sre_request.argtypes = [c_void_p, c_char_p]
        self.m_library.set_speech_engine_sre_request.restype = c_int
        ret = self.m_library.set_speech_engine_sre_request(self.m_hawkEngine, json.dumps(buffer, ensure_ascii=False).encode('utf-8'))
        return ret

    def end_enroll(self):
        buffer = {"type":"endEnroll","data":{"channel":0,"id":0,"user_id":"","text":"","start_time":-1,"end time":-1}}
        self.m_library.set_speech_engine_sre_request.argtypes = [c_void_p, c_char_p]
        self.m_library.set_speech_engine_sre_request.restype = c_int
        ret = self.m_library.set_speech_engine_sre_request(self.m_hawkEngine, json.dumps(buffer, ensure_ascii=False).encode('utf-8'))
        return ret

    def del_enroll(self, user_id):
        buffer = {"type":"deleteSpeaker","data":{"channel":0,"id":0,"user_id":user_id,"text":"","start_time":-1,"end time":-1}}
        self.m_library.set_speech_engine_sre_request.argtypes = [c_void_p, c_char_p]
        self.m_library.set_speech_engine_sre_request.restype = c_int
        ret = self.m_library.set_speech_engine_sre_request(self.m_hawkEngine, json.dumps(buffer, ensure_ascii=False).encode('utf-8'))
        return ret

    def check_enroll(self, user_id):
        buffer = {"type":"checkSpeaker","data":{"channel":0,"id":0,"user_id":user_id,"text":"","start_time":-1,"end time":-1}}
        self.m_library.set_speech_engine_sre_request.argtypes = [c_void_p, c_char_p]
        self.m_library.set_speech_engine_sre_request.restype = c_int
        ret = self.m_library.set_speech_engine_sre_request(self.m_hawkEngine, json.dumps(buffer, ensure_ascii=False).encode('utf-8'))
        return ret


    def recognize(self, start, end):
        buffer = {"type":"recognize","data":{"channel":0,"id":0,"user_id":"","text":"","start_time":start,"end time":end}}
        self.m_library.set_speech_engine_sre_request.argtypes = [c_void_p, c_char_p]
        self.m_library.set_speech_engine_sre_request.restype = c_int
        ret = self.m_library.set_speech_engine_sre_request(self.m_hawkEngine, json.dumps(buffer, ensure_ascii=False).encode('utf-8'))
        return ret
